package main

import (
    "fmt"
    "gopkg.in/yaml.v2"
    "io/ioutil"
    "log"
    "net"
    "net/http"
    "strconv"
    "strings"
    "sync"
    "time"

    "github.com/prometheus/client_golang/prometheus"
    "github.com/prometheus/client_golang/prometheus/promhttp"
)

type Config struct {
    IPs      []string `yaml:"ips"`
    Ports    []string `yaml:"ports"`
    Interval int      `yaml:"interval"`
}

var (
    portStatus = prometheus.NewGaugeVec(
        prometheus.GaugeOpts{
            Name: "port_status",
            Help: "Port status (0: closed, 1: open)",
        },
        []string{"ip", "port"},
    )
)

func init() {
    prometheus.MustRegister(portStatus)
}

func readConfig(configPath string) (*Config, error) {
    data, err := ioutil.ReadFile(configPath)
    if err != nil {
        return nil, err
    }
    var config Config
    err = yaml.Unmarshal(data, &config)
    if err != nil {
        return nil, err
    }
    return &config, nil
}

func expandRange(rangeStr string) ([]int, error) {
    var result []int
    parts := strings.Split(rangeStr, "-")
    if len(parts) == 1 {
        num, err := strconv.Atoi(parts[0])
        if err != nil {
            return nil, err
        }
        result = append(result, num)
    } else if len(parts) == 2 {
        start, err := strconv.Atoi(parts[0])
        if err != nil {
            return nil, err
        }
        end, err := strconv.Atoi(parts[1])
        if err != nil {
            return nil, err
        }
        for i := start; i <= end; i++ {
            result = append(result, i)
        }
    } else {
        return nil, fmt.Errorf("invalid range: %s", rangeStr)
    }
    return result, nil
}

func expandIPRange(ipRange string) ([]string, error) {
    var result []string
    parts := strings.Split(ipRange, "-")
    if len(parts) == 1 {
        result = append(result, parts[0])
    } else if len(parts) == 2 {
        startIP := net.ParseIP(parts[0])
        endIP := net.ParseIP(parts[1])
        for ip := startIP; !ip.Equal(endIP); ip = nextIP(ip) {
            result = append(result, ip.String())
        }
        result = append(result, endIP.String())
    } else {
        return nil, fmt.Errorf("invalid IP range: %s", ipRange)
    }
    return result, nil
}

func expandCIDR(cidr string) ([]string, error) {
    var result []string
    ip, ipnet, err := net.ParseCIDR(cidr)
    if err != nil {
        return nil, err
    }
    for ip := ip.Mask(ipnet.Mask); ipnet.Contains(ip); ip = nextIP(ip) {
        result = append(result, ip.String())
    }
    return result, nil
}

func nextIP(ip net.IP) net.IP {
    ip = ip.To4()
    for j := len(ip) - 1; j >= 0; j-- {
        ip[j]++
        if ip[j] > 0 {
            break
        }
    }
    return ip
}

func scanPort(ip string, port int, wg *sync.WaitGroup, results chan<- [3]string) {
    defer wg.Done()
    address := fmt.Sprintf("%s:%d", ip, port)
    conn, err := net.DialTimeout("tcp", address, 1*time.Second)
    status := "0"
    if err == nil {
        status = "1"
        conn.Close()
    }
    results <- [3]string{ip, strconv.Itoa(port), status}
}

func recordMetrics(results <-chan [3]string) {
    for result := range results {
        ip, port, status := result[0], result[1], result[2]
        value, _ := strconv.Atoi(status)
        portStatus.WithLabelValues(ip, port).Set(float64(value))
    }
}

func main() {
    config, err := readConfig("/mnt/config.yml")
    if err != nil {
        log.Fatalf("error: %v", err)
    }

    http.Handle("/metrics", promhttp.Handler())
    go func() {
        for {
            var ips []string
            for _, ipRange := range config.IPs {
                if strings.Contains(ipRange, "/") {
                    expandedIPs, err := expandCIDR(ipRange)
                    if err != nil {
                        log.Printf("error expanding CIDR %s: %v", ipRange, err)
                        continue
                    }
                    ips = append(ips, expandedIPs...)
                } else {
                    expandedIPs, err := expandIPRange(ipRange)
                    if err != nil {
                        log.Printf("error expanding IP range %s: %v", ipRange, err)
                        continue
                    }
                    ips = append(ips, expandedIPs...)
                }
            }

            var ports []int
            for _, portRange := range config.Ports {
                expandedPorts, err := expandRange(portRange)
                if err != nil {
                    log.Printf("error expanding port range %s: %v", portRange, err)
                    continue
                }
                ports = append(ports, expandedPorts...)
            }

            var wg sync.WaitGroup
            results := make(chan [3]string, len(ips)*len(ports))

            for _, ip := range ips {
                for _, port := range ports {
                    wg.Add(1)
                    go scanPort(ip, port, &wg, results)
                }
            }

            go func() {
                wg.Wait()
                close(results)
            }()

            recordMetrics(results)
            time.Sleep(time.Duration(config.Interval) * time.Second)
        }
    }()
    log.Fatal(http.ListenAndServe(":5000", nil))
}
